"""
GomerX 物品识别模块
用于识别和定位特定颜色的物品
"""

import cv2
import numpy as np
from modules.vision import VisionProcessor
from config.settings import VISION_CONFIG

class ObjectDetector:
    """GomerX 物品识别器"""
    
    def __init__(self, camera):
        """
        初始化物品识别器
        
        Args:
            camera: 摄像头对象
        """
        self.camera = camera
        self.vision_processor = VisionProcessor()
        self.target_color = None
        
    def set_target_color(self, color_name):
        """
        设置目标物品颜色
        
        Args:
            color_name: 颜色名称，如 "red", "blue", "green"
        """
        if color_name not in VISION_CONFIG["color_detection"]:
            raise ValueError(f"不支持的颜色: {color_name}")
            
        self.target_color = color_name
        print(f"设置目标颜色为: {color_name}")
    
    def detect_objects(self, color_name=None):
        """
        检测指定颜色的物品
        
        Args:
            color_name: 颜色名称，如果为None则使用目标颜色
            
        Returns:
            list: 检测到的物品列表，每个物品包含 (contour, position, size)
        """
        if not self.camera.is_opened:
            print("摄像头未打开")
            return []
            
        # 使用指定颜色或目标颜色
        color = color_name if color_name else self.target_color
        if not color:
            print("未指定目标颜色")
            return []
        
        # 获取摄像头图像
        frame = self.camera.read_frame()
        if frame is None:
            print("无法获取摄像头图像")
            return []
        
        # 检测颜色区域
        mask, contours = self.vision_processor.detect_color(frame, color)
        
        # 过滤可能的物体轮廓
        object_contours = self.vision_processor.filter_objects(contours)
        
        # 提取物品信息
        objects = []
        for contour in object_contours:
            # 获取物品位置和大小
            x, y, w, h, offset = self.vision_processor.get_object_position(
                contour, frame.shape[1]
            )
            
            # 计算物品面积
            area = cv2.contourArea(contour)
            
            # 存储物品信息
            objects.append({
                "contour": contour,
                "position": (x, y),
                "size": (w, h),
                "area": area,
                "offset": offset,
                "center": (x + w/2, y + h/2),
                "color": color
            })
        
        return objects
    
    def find_largest_object(self, color_name=None):
        """
        找到最大的物品
        
        Args:
            color_name: 颜色名称，如果为None则使用目标颜色
            
        Returns:
            dict: 最大物品的信息，如果没有找到则返回None
        """
        objects = self.detect_objects(color_name)
        
        if not objects:
            return None
            
        # 按面积排序，返回最大的物品
        largest_object = max(objects, key=lambda obj: obj["area"])
        return largest_object
    
    def find_closest_object(self, color_name=None):
        """
        找到最近的物品（基于图像中的位置，假设底部更近）
        
        Args:
            color_name: 颜色名称，如果为None则使用目标颜色
            
        Returns:
            dict: 最近物品的信息，如果没有找到则返回None
        """
        objects = self.detect_objects(color_name)
        
        if not objects:
            return None
            
        # 按y坐标排序，y值越大（越靠下）表示越近
        closest_object = max(objects, key=lambda obj: obj["position"][1] + obj["size"][1])
        return closest_object
    
    def is_object_in_range(self, object_info, min_area=1000, max_area=10000):
        """
        判断物品是否在合适的抓取范围内
        
        Args:
            object_info: 物品信息字典
            min_area: 最小面积阈值
            max_area: 最大面积阈值
            
        Returns:
            bool: 物品是否在合适范围内
        """
        if not object_info:
            return False
            
        area = object_info["area"]
        return min_area <= area <= max_area
    
    def get_object_distance_estimate(self, object_info):
        """
        估算物品距离（基于物品在图像中的大小）
        
        Args:
            object_info: 物品信息字典
            
        Returns:
            float: 估算的距离（相对值，越小表示越近）
        """
        if not object_info:
            return float('inf')
            
        # 使用面积的倒数作为距离估算（面积越大，距离越近）
        area = object_info["area"]
        return 1.0 / (area / 1000.0)  # 归一化处理
    
    def draw_detection_results(self, frame, objects):
        """
        在图像上绘制物品检测结果
        
        Args:
            frame: 输入图像
            objects: 检测到的物品列表
            
        Returns:
            numpy.ndarray: 绘制了检测结果的图像
        """
        result = frame.copy()
        
        for obj in objects:
            # 获取轮廓和颜色
            contour = obj["contour"]
            color_name = obj["color"]
            
            # 设置绘制颜色（BGR格式）
            if color_name == "red":
                draw_color = (0, 0, 255)
            elif color_name == "blue":
                draw_color = (255, 0, 0)
            elif color_name == "green":
                draw_color = (0, 255, 0)
            else:
                draw_color = (255, 255, 255)
            
            # 绘制轮廓
            cv2.drawContours(result, [contour], -1, draw_color, 2)
            
            # 获取边界框
            x, y = obj["position"]
            w, h = obj["size"]
            
            # 绘制边界框
            cv2.rectangle(result, (x, y), (x+w, y+h), draw_color, 2)
            
            # 添加标签
            label = f"{color_name} ({obj['area']:.0f})"
            cv2.putText(result, label, (x, y-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, draw_color, 2)
            
            # 绘制中心点
            center_x, center_y = obj["center"]
            cv2.circle(result, (int(center_x), int(center_y)), 5, draw_color, -1)
        
        return result